package middleware

import (
	"context"
	"fmt"
	"gitee.com/z930255076/go-ext/constants"
	"gitee.com/z930255076/go-ext/log"
	"gitee.com/z930255076/go-ext/response"
	"gitee.com/z930255076/go-ext/utils"
	"github.com/gofiber/fiber/v2"
	"strconv"
	"time"
)

func TimeoutMiddleware() fiber.Handler {
	return func(c *fiber.Ctx) error {
		timeSetting, _ := strconv.Atoi(utils.GetEnv("TIMEOUT", "10"))
		timeout := time.Duration(timeSetting) * time.Second
		// 创建一个带有超时的 context 并传递下去
		ctx, cancel := context.WithTimeout(c.UserContext(), timeout)
		// 取消上下文
		defer cancel()
		// 将新的 context 设置到 Fiber 的上下文中
		c.SetUserContext(ctx)

		ch := make(chan struct{}, 1)

		go func() {
			defer func() {
				if err := recover(); err != nil {
					log.PrefixLogger(fmt.Sprintf("%v", err), "TimeoutMiddleware")
				}
			}()

			if err := c.Next(); err != nil {
				log.PrefixLogger(fmt.Sprintf("%v", err), "TimeoutMiddleware")
			}

			ch <- struct{}{}
		}()

		select {
		case <-ch:
		case <-time.After(timeout):
			return response.Error(c, fiber.NewError(constants.NormalError, "请求超时"))
		}

		return nil
	}
}
